import ctypes
import ctypes.wintypes as wintypes
import sys

# Load kernel32.dll
kernel32 = ctypes.windll.LoadLibrary(r"../static/kernel32.dll")

# Define constants
PAGE_EXECUTE_READWRITE = 0x40
VIRTUAL_MEM = 0x3000
INJECT_OFFSET = 0x03
TH32CS_SNAPTHREAD = 0x00000004

# Define MEMORY_BASIC_INFORMATION structure
class MEMORY_BASIC_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("BaseAddress", ctypes.c_void_p),
        ("AllocationBase", ctypes.c_void_p),
        ("AllocationProtect", wintypes.DWORD),
        ("RegionSize", ctypes.c_size_t),
        ("State", wintypes.DWORD),
        ("Protect", wintypes.DWORD),
        ("Type", wintypes.DWORD)
    ]

# Define THREADENTRY32 structure
class THREADENTRY32(ctypes.Structure):
    _fields_ = [
        ("dwSize", wintypes.DWORD),
        ("cntUsage", wintypes.DWORD),
        ("th32ThreadID", wintypes.DWORD),
        ("th32OwnerProcessID", wintypes.DWORD),
        ("tpBasePri", ctypes.c_long),
        ("tpDeltaPri", ctypes.c_long),
        ("dwFlags", wintypes.DWORD)
    ]

# Function to search for the byte pattern in a specific memory range
def aobscan_in_range(pattern, process_handle, start_address, size):
    buffer = ctypes.create_string_buffer(size)
    bytes_read = ctypes.c_size_t()

    if kernel32.ReadProcessMemory(process_handle, start_address, buffer, size, ctypes.byref(bytes_read)):
        for i in range(size - len(pattern) + 1):
            if buffer.raw[i:i + len(pattern)] == pattern:
                return start_address + i

    return None

# Function to get all threads of a process
def get_threads(process_id):
    snapshot = kernel32.CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, process_id)
    if snapshot == -1:
        return []

    thread_entry = THREADENTRY32()
    thread_entry.dwSize = ctypes.sizeof(THREADENTRY32)
    threads = []

    if kernel32.Thread32First(snapshot, ctypes.byref(thread_entry)):
        while kernel32.Thread32Next(snapshot, ctypes.byref(thread_entry)):
            if thread_entry.th32OwnerProcessID == process_id:
                threads.append(thread_entry.th32ThreadID)

    kernel32.CloseHandle(snapshot)
    return threads

# Function to search for the byte pattern in the memory of a specific process
def aobscan(pattern, process_id):
    threads = get_threads(process_id)
    if not threads:
        print(f"No threads found for process ID {process_id}")
        return None

    process_handle = kernel32.OpenProcess(0x0010 | 0x0020 | 0x0008, False, process_id)
    if not process_handle:
        print(f"Unable to open process: {ctypes.GetLastError()}")
        return None

    mbi = MEMORY_BASIC_INFORMATION()
    base_address = ctypes.c_size_t(0)
    while kernel32.VirtualQueryEx(process_handle, base_address, ctypes.byref(mbi), ctypes.sizeof(mbi)):
        if mbi.State == 0x1000 and mbi.Protect == 0x04:
            result = aobscan_in_range(pattern, process_handle, base_address, mbi.RegionSize)
            if result:
                kernel32.CloseHandle(process_handle)
                return result
        base_address.value += mbi.RegionSize

    kernel32.CloseHandle(process_handle)
    return None

# Function to allocate memory in the target process
def alloc_memory(size, process):
    process_handle = kernel32.OpenProcess(0x0010 | 0x0020 | 0x0008, False, process)
    if not process_handle:
        print(f"Unable to open process: {ctypes.GetLastError()}")
        return None

    memory_address = kernel32.VirtualAllocEx(process_handle, 0, size, VIRTUAL_MEM, PAGE_EXECUTE_READWRITE)
    if not memory_address:
        print(f"Unable to allocate memory: {ctypes.GetLastError()}")
        return None

    kernel32.CloseHandle(process_handle)
    return memory_address

# Main function to inject code
def inject_code(process):
    pattern = b'\x01\x00\x00\xF3\x0F\x10\x46\x4C'
    inject_address = aobscan(pattern, process)

    if not inject_address:
        print("Pattern not found")
        return

    new_memory = alloc_memory(0x1000, process)

    if not new_memory:
        print("Memory allocation failed")
        return

    new_code = (
        b'\xF3\x0F\x10\x46\x50'  # movss xmm0, [rsi+50]
        b'\xF3\x0F\x11\x46\x4C'  # movss [rsi+4C], xmm0
        b'\xF3\x0F\x10\x46\x4C'  # movss xmm0, [rsi+4C]
        b'\xFF\x25' + (inject_address + 0x05).to_bytes(4, byteorder='little')  # jmp return
    )

    process_handle = kernel32.OpenProcess(0x0010 | 0x0020 | 0x0008, False, process)

    if not kernel32.WriteProcessMemory(process_handle, new_memory, new_code, len(new_code), None):
        print(f"Error writing new code: {ctypes.GetLastError()}")
        return

    jump_to_new_code = (
        b'\xE9' + (new_memory - inject_address - 0x05).to_bytes(4, byteorder='little')
    )

    if not kernel32.WriteProcessMemory(process_handle, inject_address + INJECT_OFFSET, jump_to_new_code, len(jump_to_new_code), None):
        print(f"Error writing jump to new code: {ctypes.GetLastError()}")
        return

    kernel32.CloseHandle(process_handle)
    print("Code injection successful")

# Example usage
if __name__ == "__main__":
    process_id = 9932
    inject_code(process_id)
